#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys
import numpy as np
import cv2
from collections import OrderedDict
import glob
import time
import multiprocessing
import torch

sys.path.append('.')

"""
    显示人脸
"""


def show_images(image_dir, file_names, total_image_count, dis_text='', win_name='image', wait=False):
    h_image_count = 20
    v_image_count = 10
    img_width = 96
    img_height = 96

    display_image_count = h_image_count * v_image_count
    files = []
    last_idx = len(file_names) - 1
    for idx, file_name in enumerate(file_names):
        files.append(file_name)
        file_count = len(files)
        if file_count < display_image_count and (idx < last_idx):
            continue
        imgs = []
        for file_name in files:
            file_path = os.path.join(image_dir, file_name.strip())
            img = cv2.imread(file_path)
            img = cv2.resize(img, (img_width, img_height))
            # w, h = img.shape[1], img.shape[0]
            # wscale = img_width / w
            # hscale = img_height / h
            # scale = min(wscale, hscale)
            # w = int(round(w * scale))
            # h = int(round(h * scale))
            # image = cv2.resize(img, (w, h))
            # img = np.full((img_height, img_width, image.shape[2]), fill_value=196, dtype=np.uint8)
            # x = (img_width-w)//2
            # y = (img_height-h)//2
            # img[y:y+h, x:x+w, :] = image
            imgs.append(img)
        files.clear()
        image_count = len(imgs)
        images = []
        x_off = 0
        y_off = 0
        for j, img in enumerate(imgs):
            if j % h_image_count == 0:
                x_off = 0
                if j != 0:
                    y_off += img_height
                h_images = [img]
            else:
                x_off += img_width
                h_images.append(img)
            if ((j + 1) % h_image_count == 0) or ((j + 1) == image_count):
                images.append(np.concatenate(h_images, axis=1))

        if len(images) > 1:
            w0 = images[0].shape[1]
            h0 = images[0].shape[0]
            w1 = images[-1].shape[1]
            h1 = images[-1].shape[0]
            if (w0 != w1) or (h0 != h1):
                # image = np.empty_like(images[0])
                image = np.zeros_like(images[0])
                image[0:h1, 0:w1, :] = images[-1]
                images[-1] = image
        draw_image = np.vstack(images)
        cv2.putText(draw_image, '%d/%d/%d %s' % (idx + 1, last_idx + 1, total_image_count, dis_text), (8, 24), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 255, 0), 1)
        cv2.imshow(win_name, draw_image)
        if wait:
            key = cv2.waitKey()
            if key == 27 or key == ord('q') or key == ord('Q'):
                return True
        else:
            break
    return False


def read_file_list(file_list):
    faceid_sub_files = {}
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read File Count: {}'.format(file_count))

            # if file_count == 100000:
            #     break

            # 015ff272-3768-4d4c-bebc-c14dc93065cd 477fc604-1479-4745-b8e9-9440a6ac8187 f8e4735b44724e25077106e0eaef15f0.jpg 0
            lines = line.strip().split()

            faceid = lines[0]
            sub_faceid = lines[1]
            file_name = lines[2]
            # mode = int(lines[3])

            faceid_sub = '{} {}'.format(faceid, sub_faceid)
            files = faceid_sub_files.get(faceid_sub, None)
            if files is None:
                faceid_sub_files[faceid_sub] = [file_name]
            else:
                files.append(file_name)

        print('Read File Count: {}'.format(file_count))

    return faceid_sub_files


def read_filter_list(file_list):
    faceid_subs = set()
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Filter Count: {}'.format(file_count))

            # if file_count == 100000:
            #     break

            lines = line.strip().split('_')

            faceid = lines[0]
            sub_faceid = lines[1]

            faceid_sub = '{} {}'.format(faceid, sub_faceid)
            faceid_subs.add(faceid_sub)

        print('Read File Filter: {}'.format(file_count))

    return faceid_subs


def main():
    parser = argparse.ArgumentParser(description="Display Face ReID Image")
    parser.add_argument(
        "--image_dir", default="/rootfs/media/kasim/Data1/data/VideoFaceImage", help="path to image file list", type=str
    )
    parser.add_argument(
        # "--file_list", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter5/MRMRRepFaceRecSel.txt", help="path to image file list", type=str
        "--file_list", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter6/FaceRecSel.txt", help="path to image file list", type=str
    )
    parser.add_argument(
        "--faceid_sub", default="", help="only display this face id and sub face id", type=str
    )
    parser.add_argument("--filter", action='store_true', help="filter select face")
    parser.add_argument(
        "--filter_file", default="/rootfs/media/kasim/Data/data/FaceImage/AttrDiffFace.txt", help="path to filter file", type=str
    )
    args = parser.parse_args()

    faceid_sub_files = read_file_list(args.file_list)

    if (args.faceid_sub == '') or (args.faceid_sub is None):
        if args.filter:
            filter_faceid_subs_set = read_filter_list(args.filter_file)
            for faceid_sub in filter_faceid_subs_set:
                files = faceid_sub_files.get(faceid_sub, None)
                if files is None:
                    continue
                files.sort()
                if show_images(args.image_dir, files, len(files), dis_text=faceid_sub, win_name='Face Images', wait=True):
                    break
        else:
            faceid_subs = list(faceid_sub_files.keys())
            faceid_subs.sort()
            for faceid_sub in faceid_subs:
                files = faceid_sub_files[faceid_sub]
                files.sort()
                if show_images(args.image_dir, files, len(files), dis_text=faceid_sub, win_name='Face Images', wait=True):
                    break
    else:
        files = faceid_sub_files.get(args.faceid_sub, None)
        if files is not None:
            files.sort()
            show_images(args.image_dir, files, len(files), win_name=args.faceid_sub, wait=True)

    cv2.destroyAllWindows()
    print('finish!')


if __name__ == '__main__':
    main()
